#include "common.hpp"

MPI_Win un_world_win;
void * un_world_shm;
int resv_tag;// reserved tag
int cv_tag;
int cm_tag;
int asm_tag;
int max_tag;// max tag allowed by MPI

void UnstructMG_Init()
{
    setbuf(stdout, NULL);
    assert(un_world_win == MPI_WIN_NULL);
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, & my_pid);
    if (my_pid == 0) {
        printf("Create a window\n");
        printf("Max threads %d\n", omp_get_max_threads());
    }
    MPI_Alloc_mem(sizeof(int) * 2, MPI_INFO_NULL, & un_world_shm);
    MPI_Win_create(un_world_shm, sizeof(int), sizeof(int), MPI_INFO_NULL, MPI_COMM_WORLD, & un_world_win);
    memset(un_world_shm, 0, sizeof(int) * 2);
    // MPI_Win_fence(0, un_world_win);

    void *max_tag_buf; int succ = 0;
    MPI_Comm_get_attr(MPI_COMM_WORLD, MPI_TAG_UB, &max_tag_buf, &succ); assert(succ);
    max_tag = *((int*) max_tag_buf);// MPI的最大可允许tag值
    resv_tag = max_tag - 96;
    cv_tag = resv_tag;
    cm_tag = resv_tag + 16;
    asm_tag = cm_tag + 64;
    if (my_pid == 0) printf("max tag %d => reserved %d cv %d cm %d asm %d\n", max_tag, resv_tag, cv_tag, cm_tag, asm_tag);    
}

void UnstructMG_Finalize()
{
    if (un_world_win != MPI_WIN_NULL) {
        MPI_Win_free(& un_world_win); un_world_win = MPI_WIN_NULL;
    }
    if (un_world_shm != nullptr) {
        MPI_Free_mem(un_world_shm); un_world_shm = nullptr;
    }
}

// ================================================================
//                      COO format
// ================================================================

template<typename idx_t, typename data_t, int dof>
COO<idx_t, data_t, dof>::COO() {}

template<typename idx_t, typename data_t, int dof>
COO<idx_t, data_t, dof>::COO(const COO & b) {
    this->r = b.r;
    this->c = b.c;
    for (int f = 0; f < dof*dof; f++) this->v[f] = b.v[f]; 
}

template<typename idx_t, typename data_t, int dof>
COO<idx_t, data_t, dof>& COO<idx_t, data_t, dof>::operator=(const COO & b) {
    this->r = b.r;
    this->c = b.c;
    for (int f = 0; f < dof*dof; f++) this->v[f] = b.v[f];
    return *this;
}

template<typename idx_t, typename data_t, int dof>
COO<idx_t, data_t, dof>::COO(idx_t _r, idx_t _c, data_t * _vs): r(_r), c(_c) {
    for (int f = 0; f < dof*dof; f++) this->v[f] = _vs[f];
}

template<typename idx_t, typename data_t, int dof>
bool COO<idx_t, data_t, dof>::operator == (const COO & b) const{
    return this->r == b.r && this->c == b.c;
}

template<typename idx_t, typename data_t, int dof>
bool COO<idx_t, data_t, dof>::operator < (const COO & b) const {
    if (this->r == b.r) return this->c < b.c;
    else return this->r < b.r;
}
#ifdef USE_FP16
template struct COO<int, __fp16, 1>;
template struct COO<int, __fp16, 2>;
template struct COO<int, __fp16, 3>;
template struct COO<int, __fp16, 4>;
#endif
template struct COO<int, float , 1>;
template struct COO<int, float , 2>;
template struct COO<int, float , 3>;
template struct COO<int, float , 4>;
template struct COO<int, double, 1>;
template struct COO<int, double, 2>;
template struct COO<int, double, 3>;
template struct COO<int, double, 4>;

// ================================================================
//  一些公用函数
// ================================================================

template<typename idx_t>
bool check_mask_offset(const __int128_t mask, const idx_t* offset)
{
	int num = __builtin_popcount(mask);
	int p = 0;// 指向offset位置的指针
	for (int j = 0; j < 3; j++)
	for (int i = 0; i < 3; i++)
	for (int k = 0; k < 3; k++) {
		idx_t loc = k + 3 * (i + 3 * j);
		loc = (0x1 << loc);// 左移loc位
		if (mask & loc) {// mask在该位置非零
            if (j - 1 != offset[p*3  ] ||
                i - 1 != offset[p*3+1] ||
                k - 1 != offset[p*3+2] ) {
                printf("check: mask and offset conflict! loc %d j %d i %d k %d\n",
                loc, j, i, k);
                return false;
            } 
            p++;
		}
	}
	if (p != num) return false;
	else          return true;
}
template bool check_mask_offset(const __int128_t, const int *);


template<typename idx_t, typename data_t, int dof>
void sort_indices_with_vals(const idx_t len, idx_t* inds, data_t* vals) {
    int i, j;
    constexpr int e_size = dof*dof;
    data_t buf[e_size];
    for (i = 0; i < len - 1; i++)
    for (j = 0; j < len - 1 - i; j++) {
        if (inds[j] > inds[j + 1]) {
            std::swap(inds[j], inds[j + 1]);
            for (int f = 0; f < e_size; f++) {// std::swap(vals[j], vals[j + 1]);
                buf[f] = vals[j * e_size + f];
                vals[j * e_size + f] = vals[(j + 1) * e_size + f];
                vals[(j + 1) * e_size + f] = buf[f];
            }
        }
    }
}
template void sort_indices_with_vals<int, __fp16, 1>(const int, int *, __fp16*);
template void sort_indices_with_vals<int, __fp16, 2>(const int, int *, __fp16*);
template void sort_indices_with_vals<int, __fp16, 3>(const int, int *, __fp16*);
template void sort_indices_with_vals<int, __fp16, 4>(const int, int *, __fp16*);

template void sort_indices_with_vals<int, float , 1>(const int, int *, float *);
template void sort_indices_with_vals<int, float , 2>(const int, int *, float *);
template void sort_indices_with_vals<int, float , 3>(const int, int *, float *);
template void sort_indices_with_vals<int, float , 4>(const int, int *, float *);

template void sort_indices_with_vals<int, double, 1>(const int, int *, double*);
template void sort_indices_with_vals<int, double, 2>(const int, int *, double*);
template void sort_indices_with_vals<int, double, 3>(const int, int *, double*);
template void sort_indices_with_vals<int, double, 4>(const int, int *, double*);

template<typename idx_t, typename data_t>
void sort_vals_with_indices(const idx_t len, data_t* vals, idx_t* inds)
{
  for (idx_t i = 0; i < len - 1; i++)
    for (idx_t j = 0; j < len - 1 - i; j++)
      if (vals[j] < vals[j + 1]) {
        data_t tmp_v = vals[j];
        idx_t  tmp_id= inds[j];
        vals[j] = vals[j + 1];
        inds[j] = inds[j + 1];

        vals[j + 1] = tmp_v;
        inds[j + 1] = tmp_id;
      }
}
template void sort_vals_with_indices(const int, float *, int *);
template void sort_vals_with_indices(const int, double*, int *);

template<typename k_t, typename v_t, int dof>
int uniquify(const int num, k_t * key_arr, v_t * val_arr)
{
    int k0 = 0;// 当前处理到哪个位置的元素
    int uniq_cnt = 0;// 去重后的非零元数
    constexpr int e_size = dof*dof;
    if (val_arr == nullptr) {
        while (k0 < num) {
            int k1 = k0 + 1;
            while (k1 < num && key_arr[k1] == key_arr[k0]) {
                k1 ++;
            }// 跳出循环时要么k1==num，要么遇到了一个列序号不同的点
            key_arr[uniq_cnt] = key_arr[k0];
            uniq_cnt ++;
            k0 = k1;// 移动指针
        } 
    }
    else {
        while (k0 < num) {
            int k1 = k0 + 1;
            while (k1 < num && key_arr[k1] == key_arr[k0]) {
                for (int f = 0; f < e_size; f++) val_arr[k0 * e_size + f] += val_arr[k1 * e_size + f];
                k1 ++;
            }// 跳出循环时要么k1==num，要么遇到了一个列序号不同的点
            key_arr[uniq_cnt] = key_arr[k0];
            for (int f = 0; f < e_size; f++) val_arr[uniq_cnt * e_size + f] = val_arr[k0 * e_size + f];
            uniq_cnt ++;
            k0 = k1;// 移动指针
        }
    }
    assert(k0 == num);// 跳出循环时k0==num
    return uniq_cnt;
}
template int uniquify<int, float , 1>(const int, int *, float *);
template int uniquify<int, float , 2>(const int, int *, float *);
template int uniquify<int, float , 3>(const int, int *, float *);
template int uniquify<int, float , 4>(const int, int *, float *);

template int uniquify<int, double, 1>(const int, int *, double*);
template int uniquify<int, double, 2>(const int, int *, double*);
template int uniquify<int, double, 3>(const int, int *, double*);
template int uniquify<int, double, 4>(const int, int *, double*);

template<typename idx_t, typename data_t, int dof>
int uniquify(const int num, COO<idx_t, data_t, dof> * arr)
{
    int k0 = 0;// 当前处理到哪个位置的元素
    int uniq_cnt = 0;// 去重后的非零元数
    constexpr int e_size = dof*dof;
    while (k0 < num) {
        int k1 = k0 + 1;
        while (k1 < num &&  arr[k1].r == arr[k0].r &&
                            arr[k1].c == arr[k0].c ) {
            // arr[k0].v += arr[k1].v;
            for (int f = 0; f < e_size; f++) arr[k0].v[f] += arr[k1].v[f];
            k1 ++;
        }// 跳出循环时要么k1==num，要么遇到了一个列序号不同的点
        arr[uniq_cnt] = arr[k0];
        uniq_cnt ++;
        k0 = k1;// 移动指针
    } assert(k0 == num);// 跳出循环时k0==num
    return uniq_cnt;
}
template int uniquify(const int, COO<int, float , 1> *);
template int uniquify(const int, COO<int, float , 2> *);
template int uniquify(const int, COO<int, float , 3> *);
template int uniquify(const int, COO<int, float , 4> *);
template int uniquify(const int, COO<int, double, 1> *);
template int uniquify(const int, COO<int, double, 2> *);
template int uniquify(const int, COO<int, double, 3> *);
template int uniquify(const int, COO<int, double, 4> *);

int uniq_counting(const int num, const int * arr)
{
    int k0 = 0;// 当前处理到哪个位置的元素
    int uniq_cnt = 0;// 去重后的非零元数
    while (k0 < num) {
        int k1 = k0 + 1;
        while (k1 < num &&  arr[k1] == arr[k0]) {
            k1 ++;
        }// 跳出循环时要么k1==num，要么遇到了一个列序号不同的点
        uniq_cnt ++;
        k0 = k1;// 移动指针
    } assert(k0 == num);// 跳出循环时k0==num
    return uniq_cnt;
}

double wall_time()
{
    struct timespec t;
    clock_gettime(CLOCK_MONOTONIC, &t);
    return 1. * t.tv_sec + 1.e-9 * t.tv_nsec;
}

char *get_exe_path(const char *query, char *buf, int count)
{
    int i;
    int rslt = readlink(query, buf, count - 1);
    if (rslt < 0 || (rslt >= count - 1))
    {
        return NULL;
    }
    buf[rslt] = '\0';
    for (i = rslt; i >= 0; i--)
    {
        if (buf[i] == '/')
        {
            buf[i + 1] = '\0';
            break;
        }
    }
    return buf;
}

#ifdef __x86_64__
MPI_Op MPIU_SUM_FP128;
void prod128(void* in, void* inout, int* len, MPI_Datatype* dptr)
{
  __float128* a = (__float128*) in,
            * b = (__float128*) inout;
  for (int i = 0; i < *len; i++)
    b[i] = b[i] + a[i];
}
#endif

char *print_int128(const __int128_t val, char *buf)
{
    memset(buf, -1, 128);
    __int32_t lo = -1;
    sprintf(buf +  0, "%08x", (__int32_t)((val >> 96) & lo));
    sprintf(buf +  8, "%08x", (__int32_t)((val >> 64) & lo));
    sprintf(buf + 16, "%08x", (__int32_t)((val >> 32) & lo));
    sprintf(buf + 24, "%08x", (__int32_t)( val        & lo));
    return buf;
}


void shiftl128 (unsigned int& a, unsigned int& b, unsigned int& c, unsigned int& d, size_t k)
{
    assert (k <= 128);
    if (k >= 32)
    {
        a=b;
        b=c;
        c=d;
        d=0;
        shiftl128(a,b,c,d,k-32);
    }
    else
    {
        a = (a << k) | (b >> (32-k));
        b = (b << k) | (c >> (32-k));
        c = (c << k) | (d >> (32-k));
        d = (d << k);
    }
}

void shiftr128 (unsigned int& a, unsigned int& b, unsigned int& c, unsigned int& d, size_t k)
{
    assert (k <= 128);
    if (k >= 32)
    {
        d=c;
        c=b;
        b=a;
        a=0;
        shiftr128(a,b,c,d,k-32);
    }
    else
    {
        d = (c << (32-k)) | (d >> k);
        c = (b << (32-k)) | (c >> k);
        b = (a << (32-k)) | (b >> k);
        a = (a >> k);
    }
}

__int128_t shift_left_128 (const __int128_t val, const size_t k)
{
    __int128_t ret = val;
    char * addr = (char *) (& ret);
    // 这里取地址的时候需要注意大端序和小端序
    shiftl128(*(unsigned int *)(addr + 12) , *(unsigned int *)(addr + 8), *(unsigned int *)(addr + 4), *(unsigned int *)(addr), k);
    return ret;
}

__int128_t shift_right_128 (const __int128_t val, const size_t k)
{
    __int128_t ret = val;
    char * addr = (char *) (& ret);
    // 这里取地址的时候需要注意大端序和小端序
    shiftr128(*(unsigned int *)(addr + 12) , *(unsigned int *)(addr + 8), *(unsigned int *)(addr + 4), *(unsigned int *)(addr), k);
    return ret;
}

int bitone_count(__int128_t val)
{
    __int32_t lo = -1;
    int a = __builtin_popcount( (__int32_t)((val >> 96) & lo) );
    int b = __builtin_popcount( (__int32_t)((val >> 64) & lo) );
    int c = __builtin_popcount( (__int32_t)((val >> 32) & lo) );
    int d = __builtin_popcount( (__int32_t)((val      ) & lo) );
    return a + b + c + d;
}

void thread_load_balance(int tot_work, int num_threads, int tid, int & beg, int & end)
{
    int my_work = tot_work / num_threads;
    int remain = tot_work - my_work * num_threads;
    beg = my_work * tid;
    if (tid < remain) {
        my_work ++;
        beg += tid;
    } else {
        beg += remain;
    }
    end = beg + my_work;
}

